#!/home/zhonc0a/Enthought/Canopy_64bit/User/bin/python
import matplotlib
#matplotlib.use('Agg')
import pylab
import re
import os
import sys
import argparse
import math
class bandProp():
    def __init__(self,title):
        self.title=title
        self.eneRange=''
        self.bandData=[]
        self.dossData=[]
        self.allBands=[]
        self.bandx=[]
        self.dpi=150
        self.outFileName=title
        self.bandFile='o.bands-sp1'
        self.dosFile='o.interpolated_dos'        
        self.bandRange='10000,10000'
        self.combined='1'
        self.fit=1
        self.read0weight=1
        self.intersect=1
        self.timelog='runtime'
        self.recordFile='r_em1d_ppa_HF_and_locXC_gw0'
        
    def readRecord(self):
        recordFile=open(self.recordFile)
        for line in recordFile:
            if 'Electrons' in line:
                self.HOMO=int(float(line.split(':')[1].strip('\n'))/2)
            if 'Fermi Level' in line:
                self.fermi=float(line.split(':')[1])
            if 'with scaling factors [a.u.]' in line:
                self.cella=float(line.split(':')[1].split()[0])
                self.cellb=float(line.split(':')[1].split()[1])
                self.cellc=float(line.split(':')[1].split()[2])
        recordFile.close()

    
    def readBand(self):
        bandRawData=open(self.bandFile)
        for line in bandRawData:
            if 'kx' in line and 'ky' in line and 'kz' in line:
                bandsLabel=line.split()[3:]
                lenLength=len(line.split())-1
                self.totalBand=lenLength-3
                bandDATA=[[] for i in range (self.totalBand+1)] 
            if re.search('^\s*-?[0-9]+\.[0-9]+',line) :   
                bandDATA[0].append(float(line.split()[0]))
                for i in range(3,lenLength):
                    bandDATA[i-2].append(float(line.split()[i]))
            if '# Bands Range' in line:
                self.startBand=int(line.split('|')[1])
                self.lastBand=int(line.split('|')[2]) 
        #check if kpoint data read successful, if not, change self.read0weight and try again
        if len(bandDATA[0])==0:        
            print 'ERROR: k points data still not found on '+self.bandFile
            return 0
        #print bandDATA    
        self.xpoints=len(bandDATA[0])   
        #sort bandDATA to get rid of crossing
        sortedbd=[]
        self.bandX=[i*2*math.pi/self.cella for i in bandDATA[0]]
        bd=pylab.transpose(bandDATA[1:])
        for line in bd:
            sortedbd.append(list(sorted(line)))
        bandDATA=[list(i) for i in pylab.transpose(sortedbd)]
        self.bandDATA=bandDATA

        self.bandXmin=float(self.bandX[0])
        self.bandXmax=float(self.bandX[-1])
        self.bandYmin=min(bandDATA[0])
        self.bandYmax=max(bandDATA[-1])
        
        if self.intersect==1:
            '''check for band accross and restore order
            first find points where two bands are close to each other (contact points)
            then search the the contact point to find the switch point that can make the slope of the bands most smooth
            finally sort the bands according to max of occ. band and min of virt. band
            
            '''
            for i in range(len(bandDATA)-1):
                for j in range(2,len(bandDATA[i])-2):     
                    for k in range(i+1,len(bandDATA)-1):
                        if pylab.absolute(bandDATA[i][j]-bandDATA[k][j]) < 0.1:
                            #while pylab.absolute(bandDATA[i][j]-bandDATA[k][j]) > pylab.absolute(bandDATA[i][j+1]-bandDATA[k][j+1]):
                            #    j=j+1
                            if pylab.absolute(2*bandDATA[i][j]-bandDATA[i][j-1]-bandDATA[i][j+1]) > pylab.absolute(2*bandDATA[k][j]-bandDATA[i][j-1]-bandDATA[k][j+1]):
                                while pylab.absolute(2*bandDATA[k][j]-bandDATA[i][j-1]-bandDATA[k][j+1]) > pylab.absolute(2*bandDATA[k][j+1]-bandDATA[i][j]-bandDATA[k][j+2]) and j+3 < len(bandDATA[k]):
                                    j=j+1
                                topl=bandDATA[k][:j]
                                topr=bandDATA[k][j:]
                                btnl=bandDATA[i][:j]
                                btnr=bandDATA[i][j:]
                                bandDATA[i]=btnl+topr
                                bandDATA[k]=topl+btnr
            sortedOcc=sorted(bandDATA[:self.HOMO],key=lambda kpoint: max(kpoint))
            sortedVirt=sorted(bandDATA[self.HOMO:],key=lambda kpoint: min(kpoint))
            bandDATA=sortedOcc+sortedVirt  
        self.allBands=bandDATA

        
    #def readTime(self):
    #    timelog=open(self.timelog,'r')
    #    firstLine=timelog.readline()
    #    rawtime=firstLine.split()[2].replace('elapsed','')
    #    second=float(rawtime.split(':')[-1])
    #    try:
    #        minute=float(rawtime.split(':')[-2])*60
    #    except IndexError:
    #        minute=0
    #    try:
    #        hour=float(rawtime.split(':')[-3])*3600
    #    except IndexError:
    #        hour=0
    #    self.runtime=second+minute+hour
        
    def readDOSS(self):
        dossRawData=open(self.dosFile)
        for line in dossRawData:
            if re.search('-?\d+\.\d+\s+-?\d+\.\d+',line) and len(line.split())==2:
                self.dossData.append(line.split())
        self.npdos=len(self.dossData[0])-1
        self.dossX=pylab.array(map(float,pylab.array(self.dossData).transpose()[0] ))
        #self.pdossY=[]
        #try:
        #    if max(map(int,self.pdosColumn)) > self.npdos-1:
        #        print 'number of input pdos column '+str(self.pdosColumn)+' exceed number of pdos in file: '+str(self.npdos-1)
        #except ValueError:
        #    pass
        #for i in range(1,self.npdos):
        #    if str(i) in self.pdosColumn:
        #        self.pdossY.append(pylab.array(map(float,pylab.array(self.dossData).transpose()[i] )))
        self.dossY=pylab.array(map(float,pylab.array(self.dossData).transpose()[-1] ))
        #if self.gaussian>0:
        #    self.FWHM=((self.dossX[0]-self.dossX[-1])/len(self.dossX)*self.gaussian)*27.211396*2.355
        #    WdossY=gaussian_filter1d(self.dossY,self.gaussian,truncate=100.0,mode='constant')
        #    WpdossY=[]
        #    for i in range(len(self.pdossY)):
        #        WpdossY.append(gaussian_filter1d(self.pdossY[i],self.gaussian,truncate=100.0,mode='constant'))
        #    self.dossY=WdossY
        #    self.pdossY=WpdossY
        self.dossYmin=self.dossY.min()
        self.dossYmax=self.dossY.max()
        
    def propGen(self):
        HOMO=self.HOMO-self.startBand
        self.HOMOband=pylab.array(self.allBands[HOMO])/27.211396
        self.LUMOband=pylab.array(self.allBands[HOMO+1])/27.211396
        HOMOmax=self.HOMOband.max()
        HOMOmin=self.HOMOband.min()
        LUMOmax=self.LUMOband.max()
        LUMOmin=self.LUMOband.min()
        self.HOMOmax=HOMOmax*27.211396
        self.LUMOmin=LUMOmin*27.211396
        self.HOMObandWidth=(HOMOmax-HOMOmin)*27.211396
        self.LUMObandWidth=(LUMOmax-LUMOmin)*27.211396
        self.bandGap=(LUMOmin-HOMOmax)*27.211396
        if self.fit==1:
            #calculate effective mass
            '''1 cut the band into n segment
            the segment involve the HOMO max or LUMO min will be used for fitting
            use parabole to fit the curve, calculate effective mass, R square and mean absolute error
            if the R squre > 0.9995 and MAE < 0.0001. the effective mass will be adopted'''
            fitErrorHOMO=[]
            fitRHOMO=[]
            fitErrorLUMO=[]
            fitRLUMO=[]
            massHole=[]
            massElectron=[]
            fitX=[]
            totalPoints=len(self.HOMOband)
            for n in range(2,self.xpoints/5+2):
                maxIndex=self.HOMOband.argmax()
                if maxIndex<len(self.HOMOband)/n:
                    partHOMO=self.HOMOband[:len(self.HOMOband)/n]
                    partXH=self.bandX[:len(self.HOMOband)/n]
                elif maxIndex>len(self.HOMOband)/n*(n-1):
                    partHOMO=self.HOMOband[-len(self.HOMOband)/n:]
                    partXH=self.bandX[-len(self.HOMOband)/n:]
                else:
                    partHOMO=self.HOMOband[maxIndex-len(self.HOMOband)/(2*n):maxIndex+len(self.HOMOband)/(2*n)]
                    partXH=self.bandX[maxIndex-len(self.HOMOband)/(2*n):maxIndex+len(self.HOMOband)/(2*n)]
                minIndex=self.LUMOband.argmin()
                if minIndex<len(self.LUMOband)/n:
                    partLUMO=self.LUMOband[:len(self.LUMOband)/n]
                    partXL=self.bandX[:len(self.LUMOband)/n]
                elif minIndex>len(self.LUMOband)/n*(n-1):
                    partLUMO=self.LUMOband[-len(self.LUMOband)/n:]
                    partXL=self.bandX[-len(self.LUMOband)/n:]
                else:
                    partLUMO=self.LUMOband[minIndex-len(self.LUMOband)/(2*n):minIndex+len(self.LUMOband)/(2*n)]
                    partXL=self.bandX[minIndex-len(self.LUMOband)/(2*n):minIndex+len(self.LUMOband)/(2*n)]
                coefficients=pylab.polyfit(partXH,partHOMO,2)
                polynomial = pylab.poly1d(coefficients)
                d2=polynomial.deriv().deriv()
                massHole.append(d2(self.bandX[maxIndex])**-1)
                fitHOMO=polynomial(partXH)
                HOMObar=pylab.sum(partHOMO)/len(partHOMO)
                ssreg=pylab.sum((fitHOMO-HOMObar)**2)
                sstot=pylab.sum((partHOMO-HOMObar)**2)
                fitErrorHOMO.append(pylab.sum(pylab.absolute(fitHOMO-partHOMO))/len(partHOMO))
                fitRHOMO.append(ssreg/sstot)
                coefficients=pylab.polyfit(partXL,partLUMO,2)
                polynomial = pylab.poly1d(coefficients)
                d2=polynomial.deriv().deriv()
                fitLUMO=polynomial(partXL)
                LUMObar=pylab.sum(partLUMO)/len(partLUMO)
                ssreg=pylab.sum((fitLUMO-LUMObar)**2)
                sstot=pylab.sum((partLUMO-LUMObar)**2)
                fitErrorLUMO.append(pylab.sum(pylab.absolute(fitLUMO-partLUMO))/len(partLUMO))
                fitRLUMO.append(ssreg/sstot)
                massElectron.append(d2(self.bandX[minIndex])**-1)
                fitX.append(n)  
            self.holeMass='---'
            self.electronMass='---'
            for n in range(self.xpoints/5):
                if fitRHOMO[n]>0.999 and fitErrorHOMO[n]<0.0005:
                    self.holeMass=massHole[n]
                    break
            for n in range(self.xpoints/5):
                if fitRLUMO[n]>0.999 and fitErrorLUMO[n]<0.0005:
                    self.electronMass=massElectron[n]
                    break
            if self.holeMass != '---' or self.electronMass != '---':
                pylab.figure('massFit')
                pylab.subplot(3,1,1)
                pylab.plot(fitX,massHole,'b-',fitX,massElectron,'r-')
                pylab.ylabel('Effective Mass')
                pylab.subplot(3,1,2)
                pylab.plot(fitX,fitRHOMO,'b-',fitX,fitRLUMO,'r-')
                pylab.ylabel('R squre')
                pylab.subplot(3,1,3)
                pylab.plot(fitX,fitErrorHOMO,'b-',fitX,fitErrorLUMO,'r-')
                pylab.ylabel('Mean Absolute Error')
                pylab.suptitle('Number of points is '+str(totalPoints))
                pylab.savefig(self.title+'_massFit.png',dpi=self.dpi,bbox_inches='tight',format='png')
        else:
            self.holeMass='---'
            self.electronMass='---'
        
#test=bandProp('PA')
#test.readBand()    
#for band in test.allBands:
#    pylab.plot(test.bandX,band)
#pylab.show()   
#test.propGen()
#print test.holeMass
#print test.electronMass
#print test.HOMObandWidth
#print test.LUMObandWidth

#        
    def plotBand(self):
        downRange=int(self.bandRange.split(',')[0])
        upRange=int(self.bandRange.split(',')[1])
        HOMO=self.HOMO-self.startBand
        startBand=HOMO-downRange
        lastBand=HOMO+upRange
        if startBand < 0:
            startBand=0
        if lastBand > self.totalBand:
            lastBand=self.totalBand-1
        #print startBand
        #print lastBand
        #occBand=self.HOMO-1
        #virBand=lastBand-self.HOMO

        pylab.figure('Band Structure')
        if self.combined=='1':
            pylab.subplot(121)
        for i in range(startBand,lastBand+1):
            if i <= HOMO:
                pylab.plot(self.bandX,pylab.array(self.allBands[i]),'b-')
            if i > HOMO:
                pylab.plot(self.bandX,pylab.array(self.allBands[i]),'r-') 
        bandYmin=pylab.array(self.allBands[startBand]).min()
        bandYmax=pylab.array(self.allBands[lastBand]).max()
        bandYmargin=(bandYmax-bandYmin)*0.05
        pylab.axis([self.bandXmin,self.bandXmax,bandYmin-bandYmargin,bandYmax+bandYmargin])
        pylab.ylabel('Energy (eV)')
        pylab.xlabel('K (Bohr-1)')
        if self.combined=='0':
            pylab.savefig(self.title+'_bands.png',dpi=self.dpi,bbox_inches='tight',format='png')
            pylab.savefig(self.title+'_bands.svg',bbox_inches='tight',format='svg')
        if self.combined=='1':
            pylab.subplot(122)
        if self.combined=='0':
            pylab.figure('Density of State')
        pylab.plot(self.dossY,self.dossX,'b-')
        pylab.axis([self.dossYmin,self.dossYmax,bandYmin-bandYmargin,bandYmax+bandYmargin])
        pylab.xlabel('State Density')
        if self.combined=='1':
            pylab.ylabel('Energy (eV)')
        if self.combined=='1':
            
            pylab.savefig(self.title+'_band-doss.png',dpi=self.dpi,bbox_inches='tight',format='png')
            pylab.savefig(self.title+'_band-doss.svg',bbox_inches='tight',format='svg')
        if self.combined=='0':
            pylab.savefig(self.title+'_doss.png',dpi=self.dpi,bbox_inches='tight',format='png')
            pylab.savefig(self.title+'_doss.svg',bbox_inches='tight',format='svg')
#    
parser=argparse.ArgumentParser(description='Generate band structure figure and corresponding data')
parser.add_argument("-b","--band",metavar='bandFile',help='Band structure file generate by crystal14, default:o.bands-sp1',type=str)
parser.add_argument("-d","--doss",metavar='dossFile',help='Dstructure file generate by crystal14, default:o.interpolated_dos',type=str)
parser.add_argument("-r","--range",help='range for draw band structure, e.g. 1,2 means draw form HOMO-1 to LUMO+1, default:all the band',type=str)
parser.add_argument("-c","--combined",help='draw Band structure and Doss seperately or combined,1=combined, 0=seperately, default:1',type=str)
parser.add_argument("-m","--mass",help='fit the band structure to calculate effective mass,1=calculate, other=not calculate, default:1',default=0,type=int)
#parser.add_argument("-w","--bandWeight",help='select the type of band to draw,1=bands with 0 weight, other=bands with non-zeor weight, default:1',default=1,type=int)
parser.add_argument("-i","--intersect",help='whether to allow band to intersect ,1=intersect, other=not intersect, default:0',default=0,type=int)
parser.add_argument("-v","--value",help='the value to draw in the converge graph. (gap,home,lumo,hw,lw). default:None',default='',type=str)
parser.add_argument("-p","--plot",help='plot the band structure or not. 1=plot, 0=not plot. default:0',default='0',type=str)
#parser.add_argument("-s","--size",help='the dpi of the figure,0=combined, 1=seperately, default:1',type=str)

parser.add_argument("targetDIR",nargs='*')
args=parser.parse_args()

def genColorGrad(step):
    if step==1:
        return [(1.0,0.0,0.0)]
    if step==2:
        return [(1.0,0.0,0.0),(0.0,0.0,1.0)]
    rg=step/2+1
    gb=step/2+1
    if rg+gb>step+1:
        gb=gb-1
    rglist=[]
    gblist=[]

    for i in range(rg):
        rglist.append((1-i/(rg-1.0),i/(rg-1.0),0.0))
    for i in range(gb):
        gblist.append((0.0,1.0-i/(gb-1.0),i/(gb-1.0)))
    rglist.pop()
    return rglist+gblist





if args.targetDIR and args.value:
    originDIR=os.getcwd()
    print "%-35s%-12s%-12s%-12s%-15s%-15s%-12s%-12s" % ('file name','HOMO Width','LUMO Width','Band Gap','Hole Mass','Electron Mass','HOMO','LUMO')
    scanedP={'K':[],'B':[],'W':[],'E':[],'G':[],'e':[]}
    scanedValue={'K':{},'B':{},'W':{},'E':{},'G':{},'e':{}}
    scanedTime={'K':{},'B':{},'W':{},'E':{},'G':{},'e':{}}
    for targetDIR in args.targetDIR:
        if not os.path.exists(targetDIR):
            print "Error!!! Target Dierctory "+targetDIR+" do not exist!"
            sys.exit()
        targetName=''.join(re.split('[/]',targetDIR)[-1]).split('.')[0]
        pV=targetName.split('_')[0]
        pm=targetName.split('_')[0][0]
        pB='_'.join(targetName.split('_')[1:])  
        if pV not in scanedP[pm]:
            scanedP[pm].append(pV)
        if pB not in scanedValue[pm]:
            scanedValue[pm][pB]=[]
        if pB not in scanedTime[pm]:
            scanedTime[pm][pB]=[]
        bandgen=bandProp(targetName)
        bandgen.fit=args.mass
        bandgen.read0weight=args.bandWeight
        bandgen.intersect=args.intersect
        if args.range:
            bandgen.bandRange=args.range
        if args.combined:
            bandgen.combined=args.combined
        os.chdir(targetDIR)
        bandgen.readRecord()
        if bandgen.readBand()!=0:
            
            bandgen.readDOSS()
            bandgen.propGen()
            bandgen.readTime()
            scanedTime[pm][pB].append(bandgen.runtime)
        
        
        #if bandgen.read0weight==1:
        #    print str(bandgen.xpoints)+' k points with zero weight'
        #else:
        #    print str(bandgen.xpoints)+' k points with non-zero weight:'
            if args.value.lower()=='gap':
                scanedValue[pm][pB].append(bandgen.bandGap)
                graphTitle='gap'
            elif args.value.lower()=='homo':
                scanedValue[pm][pB].append(bandgen.HOMOmax)
                graphTitle='HOMO'
            elif args.value.lower()=='lumo':
                scanedValue[pm][pB].append(bandgen.LUMOmin)
                graphTitle='LUMO'
            elif args.value.lower()=='hw':
                scanedValue[pm][pB].append(bandgen.HOMObandWidth)
                graphTitle='HOMO width'
            elif args.value.lower()=='lw':
                scanedValue[pm][pB].append(bandgen.LUMObandWidth)
                graphTitle='LUMO width'
            else:
                print 'value "'+args.value+'" not recognized' 
                sys.exit()       
            print "%-35s%-12.6f%-12.6f%-12.6f%-15.6s%-15.6s%-12.6f%-12.6f" % (bandgen.title,bandgen.HOMObandWidth,bandgen.LUMObandWidth,bandgen.bandGap,bandgen.holeMass,bandgen.electronMass,bandgen.HOMOmax,bandgen.LUMOmin)
            if args.plot=='1':
                if args.combined:
                    bandgen.combined=args.combined
                bandgen.plotBand()
        else:
            print "%-35s%-s%" % (bandgen.title,"ERROR: k points data still not found on EIGENVAL.")
        os.chdir(originDIR)
    #print scanedP
    #print scanedValue
    
    
    
    for i in scanedP:
        if len(scanedP[i]) != 0:
            pylab.figure(i)
            pylab.clf()
            x=range(len(scanedP[i]))
            pylab.xticks(x,scanedP[i])
            lineColors=genColorGrad(len(scanedValue[i]))
            k=0
            for j in scanedValue[i]:
                pylab.plot(x,scanedValue[i][j],color=lineColors[k],label=j)
                pylab.ylabel(graphTitle)
                k=k+1
            pylab.legend()
            pylab.savefig(bandgen.title+'_conver.png',dpi=bandgen.dpi,bbox_inches='tight',format='png')
            pylab.savefig(bandgen.title+'_conver.svg',bbox_inches='tight',format='svg')
            pylab.clf()
            pylab.xticks(x,scanedP[i])
            k=0
            for j in scanedTime[i]:
                pylab.plot(x,scanedTime[i][j],color=lineColors[k],label=j)
                pylab.ylabel('Time (second)')
                k=k+1
            pylab.legend()
            pylab.savefig(bandgen.title+'_time.png',dpi=bandgen.dpi,bbox_inches='tight',format='png')
            pylab.savefig(bandgen.title+'_time.svg',bbox_inches='tight',format='svg')
            
elif args.targetDIR:
    originDIR=os.getcwd()
    print "%-35s%-12s%-12s%-12s%-15s%-15s%-12s%-12s" % ('file name','HOMO Width','LUMO Width','Band Gap','Hole Mass','Electron Mass','HOMO','LUMO')
    for targetDIR in args.targetDIR:
        if not os.path.exists(targetDIR):
            print "Error!!! Target Dierctory "+targetDIR+" do not exist!"
            sys.exit()       
        targetName=''.join(re.split('[/]',targetDIR)[-1])    
        bandgen=bandProp(targetName)
        bandgen.fit=args.mass
        bandgen.read0weight=args.bandWeight
        bandgen.intersect=args.intersect
        if args.range:
            bandgen.bandRange=args.range
        if args.combined:
            bandgen.combined=args.combined
        os.chdir(targetDIR)
        bandgen.readRecord()
        if bandgen.readBand()!=0:
        #bandgen.readDOSS()
            bandgen.propGen()     
            print "%-35s%-12.6f%-12.6f%-12.6f%-15.6s%-15.6s%-12.6f%-12.6f" % (bandgen.title,bandgen.HOMObandWidth,bandgen.LUMObandWidth,bandgen.bandGap,bandgen.holeMass,bandgen.electronMass,bandgen.HOMOmax,bandgen.LUMOmin)
            if args.plot=='1':
                if args.combined:
                    bandgen.combined=args.combined
                bandgen.plotBand()
        else:        
            print "%-35s%-s%" % (bandgen.title,"ERROR: k points data still not found on EIGENVAL.")
        os.chdir(originDIR)      
                
else:
    bandgen=bandProp('PA')
    bandgen.fit=args.mass
    #bandgen.read0weight=args.bandWeight
    bandgen.intersect=args.intersect
    #if args.band:
    #    bandgen.bandFileName=args.band
    #    if args.doss:
    #        bandgen.dossFileName=args.doss
    if args.range:
        bandgen.bandRange=args.range
    if args.combined:
        bandgen.combined=args.combined
    bandgen.readRecord()
    bandgen.readBand()
    bandgen.readDOSS()
    bandgen.propGen()
    if bandgen.read0weight==1:
        print str(bandgen.xpoints)+' k points with zero weight'
    else:
        print str(bandgen.xpoints)+' k points with non-zero weight:'
    print "%-35s%-12s%-12s%-12s%-15s%-15s%-12s%-12s" % ('file name','HOMO Width','LUMO Width','Band Gap','Hole Mass','Electron Mass','HOMO','LUMO')
    print "%-35s%-12.6f%-12.6f%-12.6f%-15.6s%-15.6s%-12.6f%-12.6f" % (bandgen.title,bandgen.HOMObandWidth,bandgen.LUMObandWidth,bandgen.bandGap,bandgen.holeMass,bandgen.electronMass,bandgen.HOMOmax,bandgen.LUMOmin)
    if args.plot=='1':
        if args.combined:
            bandgen.combined=args.combined
        bandgen.plotBand()
    #    
#    
    
    #if last_max==pylab.array(map(float,line)).min()
    #last_max=pylab.array(map(float,line)).max()
    #last_min=pylab.array(map(float,line)).min()
#    





